package rag

import (
	"bytes"
	"context"
	"encoding/csv"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"quipus/internal/llm"
	"quipus/internal/rag/storage"
	ragStore "quipus/internal/rag/storage"
	"sort"

	cpb "quipus/proto/gen/common"

	"strings"
	"text/template"
)

type KeyWord struct {
	LowLevelKeywords  []string `json:"low_level_keywords,omitempty"`
	HighLevelKeywords []string `json:"high_level_keywords,omitempty"`
}

type RagContent struct {
	Entities      string `json:"entities,omitempty"`
	Relationships string `json:"relationships,omitempty"`
	Sources       string `json:"sources,omitempty"`
}

type EntityList []ragStore.Entity
type BaseEntityList []ragStore.BaseEntity

type EntityReverseList []ragStore.ReverseEntity

type RelationList []ragStore.Relation
type BaseRelationList []ragStore.BaseRelation

type RelationEntityList []ragStore.RelationEntity

func (b *RagBuilder) naiveQueryLLMPrompt(content string, docIds []string, useDefault bool) (string, error) {
	txt := ""
	chunks := []*ragStore.ChunkVector{}
	for _, docId := range docIds {
		cks, err := b.vectordb.QueryChunk(content, docId, b.doEmbedding)
		if err != nil {
			return "", err
		}
		chunks = append(chunks, cks...)
	}

	//rerank
	sort.Slice(chunks, func(i, j int) bool {
		// 距离越小越好
		return chunks[i].Score < chunks[j].Score
	})

	basicScore := chunks[0].Score

	tokensLen := 0
	for _, chunk := range chunks {
		if txt == "" {
			txt = chunk.Content
			continue
		}

		//若选出来的值相关性偏差较大的时候，舍弃
		if chunk.Score-basicScore > b.ScoreBias {
			break
		}

		tokens := b.Tkm.Encode(chunk.Content, nil, nil)
		tokensLen += len(tokens)
		if len(tokens) < b.MaxTokens {
			txt += chunk.Content
		} else {
			break
		}
	}

	prompt, err := b.buildNaiveSystemPromt(&txt, &content, useDefault)
	if err != nil {
		return "", err
	}

	return prompt, nil
}

func (b *RagBuilder) NaiveQueryLLM(content string, docId string) (string, error) {
	prompt, err := b.naiveQueryLLMPrompt(content, []string{docId}, true)
	if err != nil {
		return "", err
	}
	ctx := context.Background()
	slog.Info("query llm  ", slog.String("prompt:", prompt))
	return b.QueryLLM(ctx, content, prompt)
}

func (b *RagBuilder) QueryLLMRagPrompt(content string, docIds []string, useDefault bool, ragMode cpb.RagMode) (string, error) {
	switch ragMode {
	case cpb.RagMode_NAIVE:
		return b.naiveQueryLLMPrompt(content, docIds, useDefault)
	case cpb.RagMode_KG:
		ctx := context.Background()
		return b.QueryRagPrompt(ctx, content, docIds, useDefault)
	default:
		return b.naiveQueryLLMPrompt(content, docIds, useDefault)
	}
}

func (b *RagBuilder) QueryRag(content string, docId string) (string, error) {
	if len(content) < 50 {
		return b.NaiveQueryLLM(content, docId)
	}

	ctx := context.Background()
	prompt, err := b.QueryRagPrompt(ctx, content, []string{docId}, true)
	if err != nil {
		return "", err
	}
	return b.QueryLLM(ctx, content, prompt)
}

func (b *RagBuilder) QueryRagPrompt(ctx context.Context, content string, docIds []string, useDefault bool) (string, error) {
	kwPromt, err := b.buildKeyWordPromt(content)
	if err != nil {
		return "", err
	}

	var kw KeyWord

	decode := func(msg string) error {
		return json.Unmarshal([]byte(msg), kw)
	}
	err = b.UnmarshalLLMResponse(ctx, kwPromt, decode)
	if err != nil {
		return "", err
	}

	var entityList *BaseEntityList
	var relationList *BaseRelationList
	if len(kw.LowLevelKeywords) > 0 {
		for _, docId := range docIds {
			lowEntityList, lowRelationList, err := b.buildLowPrompt(strings.Join(kw.LowLevelKeywords, ","), docId)
			if err != nil {
				return "", nil
			}
			if entityList == nil {
				entityList = lowEntityList
			} else {
				mergeEntityList(entityList, lowEntityList)
			}

			if relationList == nil {
				relationList = lowRelationList
			} else {
				mergeRelationList(relationList, lowRelationList)
			}
		}

	}

	if len(kw.HighLevelKeywords) > 0 {
		for _, docId := range docIds {
			highEntityList, highRelationList, err := b.buildHighPrompt(strings.Join(kw.HighLevelKeywords, ","), docId)
			if err != nil {
				return "", nil
			}

			if entityList == nil {
				entityList = highEntityList
			} else {
				mergeEntityList(entityList, highEntityList)
			}

			if relationList == nil {
				relationList = highRelationList
			} else {
				mergeRelationList(relationList, highRelationList)
			}
		}

	}
	if entityList != nil {
		sort.Stable(entityList)
	}
	if relationList != nil {
		sort.Stable(relationList)
	}

	ragContent, err := b.buildRagContentPrompt(entityList, relationList, docIds)
	return b.buildSystemPromt(ragContent, useDefault)
}

type EntitySet struct {
	items map[string]*storage.BaseEntity // 使用某个唯一字段作为键
}

func (s *EntitySet) init() {
	s.items = make(map[string]*storage.BaseEntity)
}
func (s *EntitySet) Add(entity storage.BaseEntity) {
	if s.items == nil {
		s.init()
	}

	// 假设 entity 有一个 ID 或其他唯一标识符字段
	e := s.items[entity.ChunkId]
	if e != nil {
		return
	}
	s.items[entity.ChunkId] = &entity
}
func (s *EntitySet) Cardinality() int {
	return len(s.items)
}

func (s *EntitySet) Each(f func(string) bool) {
	for k, _ := range s.items {

		if !f(k) {
			break
		}
	}
}

func (b *RagBuilder) buildRagContentPrompt(entityList *BaseEntityList, relationList *BaseRelationList, docIds []string) (*RagContent, error) {
	docMap := make(map[string]*EntitySet)
	if entityList != nil {
		for _, entity := range *entityList {
			chunkSet := docMap[entity.DocId]
			if chunkSet == nil {
				chunkSet = &EntitySet{}
				docMap[entity.DocId] = chunkSet
			}
			chunkSet.Add(entity)
		}
	}

	var chunks []*ragStore.ChunkVector
	if len(docMap) < 1 {
		return nil, nil
	}

	for docId, cksets := range docMap {
		cks, err := b.getChunkVector(cksets, docId)
		if err != nil {
			return nil, err
		}
		chunks = append(chunks, cks...)
	}

	computeToken := func(txt string) string {
		tokens := b.Tkm.Encode(txt, nil, nil)
		if len(tokens) < b.MaxTokens {
			return txt
		}
		return b.Tkm.Decode(tokens[:b.MaxTokens-1])
	}

	ragContent := &RagContent{
		Entities:      computeToken(EntityToCSV(entityList)),
		Relationships: computeToken(RelationToCSV(relationList)),
		Sources:       computeToken(ChunkToCSV(chunks)),
	}

	return ragContent, nil
}

func mergeEntityList(entityList *BaseEntityList, entityList2 *BaseEntityList) {
	if entityList == nil {
		entityList = entityList2
	} else {
		var entityTmpList BaseEntityList
		for _, e2 := range *entityList2 {
			merge := true
			for _, e := range *entityList {
				if e.Uid == e2.Uid {
					merge = false
					break
				}
			}

			if merge {
				entityTmpList = append(entityTmpList, e2)
			}

		}
		*entityList = append(*entityList, entityTmpList...)
	}
}

func mergeRelationList(relationList *BaseRelationList, relationList2 *BaseRelationList) {
	if relationList == nil {
		relationList = relationList2
	} else {
		var relationTmpList BaseRelationList
		for _, r2 := range *relationList2 {
			merge := true
			for _, r := range *relationList {
				if r.Uid == r2.Uid {
					merge = false
					break
				}
			}
			if merge {
				relationTmpList = append(relationTmpList, r2)
			}

		}
		*relationList = append(*relationList, relationTmpList...)
	}
}

func (l BaseEntityList) Len() int {
	return len(l)
}

func (l BaseEntityList) Less(i, j int) bool {
	return l[i].Rank > l[j].Rank
}

func (l BaseEntityList) Swap(i, j int) {
	l[i], l[j] = l[j], l[i]
}

func (l BaseRelationList) Len() int {
	return len(l)
}

func (l BaseRelationList) Less(i, j int) bool {
	return l[i].Rank > l[j].Rank
}

func (l BaseRelationList) Swap(i, j int) {
	l[i], l[j] = l[j], l[i]
}

func EntityToCSV(records *BaseEntityList) string {
	// 创建一个csv.Writer实例
	var ecsv strings.Builder

	ecsvWriter := csv.NewWriter(&ecsv)

	// 将结构体字段写入CSV标题行
	if err := ecsvWriter.Write([]string{"Name", "Content", "Rank"}); err != nil {
		panic(err)
	}

	// 将每个结构体实例转换为CSV记录行并写入
	for _, record := range *records {
		row := []string{record.Name, record.EntityDesc, fmt.Sprintf("%d", record.Rank)}
		if err := ecsvWriter.Write(row); err != nil {
			panic(err)
		}
	}

	ecsvWriter.Flush()
	if err := ecsvWriter.Error(); err != nil {
		panic(err)
	}

	return ecsv.String()
}

func RelationToCSV(records *BaseRelationList) string {
	// 创建一个csv.Writer实例
	var csvContent strings.Builder
	csvWriter := csv.NewWriter(&csvContent)

	// 将结构体字段写入CSV标题行
	if err := csvWriter.Write([]string{"SrcName", "TgtName", "Content", "KeyWord", "Rank"}); err != nil {
		panic(err)
	}

	// 将每个结构体实例转换为CSV记录行并写入
	for _, record := range *records {
		row := []string{record.SrcName, record.TgtName, record.RelDesc, record.RelKeyWords, fmt.Sprintf("%d", record.Rank)}
		if err := csvWriter.Write(row); err != nil {
			panic(err)
		}
	}

	csvWriter.Flush()
	if err := csvWriter.Error(); err != nil {
		panic(err)
	}

	return csvContent.String()
}

func ChunkToCSV(records []*ragStore.ChunkVector) string {
	// 创建一个csv.Writer实例
	var csvContent strings.Builder
	csvWriter := csv.NewWriter(&csvContent)

	// 将结构体字段写入CSV标题行
	if err := csvWriter.Write([]string{"Seq", "Content"}); err != nil {
		panic(err)
	}

	// 将每个结构体实例转换为CSV记录行并写入
	for i, record := range records {
		row := []string{fmt.Sprintf("%d", i), record.Content}
		if err := csvWriter.Write(row); err != nil {
			panic(err)
		}
	}

	csvWriter.Flush()
	if err := csvWriter.Error(); err != nil {
		panic(err)
	}

	return csvContent.String()
}

func (b *RagBuilder) buildLowPrompt(content string, docId string) (*BaseEntityList, *BaseRelationList, error) {
	docVectors, err := b.queryEntityVector(content, docId)
	if err != nil {
		return nil, nil, err
	}

	entityList, err := b.queryEntityGraph(docVectors, docId)
	entityReverseList, err := b.queryEntityReverseGraph(docVectors, docId)

	BaseEntityList := make(BaseEntityList, 0)
	BaseRelationList := make(BaseRelationList, 0)

	// chunkSet := mapset.NewSet()
	for _, entity := range *entityList {
		// chunkSet.Add(entity.ChunkId)
		BaseEntityList = append(BaseEntityList, entity.BaseEntity)
		if entity.Owns != nil {
			for _, rel := range entity.Owns {
				BaseRelationList = append(BaseRelationList, rel.BaseRelation)
				for _, tgtEntity := range rel.Owes {
					BaseEntityList = append(BaseEntityList, tgtEntity.BaseEntity)
				}
			}
		}
	}

	for _, entity := range *entityReverseList {
		BaseEntityList = append(BaseEntityList, entity.BaseEntity)
		if entity.Owes != nil {
			for _, rel := range entity.Owes {
				BaseRelationList = append(BaseRelationList, rel.BaseRelation)
				for _, tgtEntity := range rel.Owes {
					BaseEntityList = append(BaseEntityList, tgtEntity.BaseEntity)
				}
			}
		}

	}

	return &BaseEntityList, &BaseRelationList, nil
}

func (b *RagBuilder) buildHighPrompt(content string, docId string) (*BaseEntityList, *BaseRelationList, error) {
	docVectors, err := b.queryRelationVector(content, docId)
	if err != nil {
		return nil, nil, err
	}

	BaseEntityList := make(BaseEntityList, 0)
	BaseRelationList := make(BaseRelationList, 0)

	relationList, err := b.queryRelationGraph(docVectors, docId)
	if err != nil {
		return nil, nil, err
	}
	// chunkSet := mapset.NewSet()
	if relationList != nil {
		for _, re := range *relationList {
			re.Rank = len(re.Owes) + len(re.Owns)
			BaseRelationList = append(BaseRelationList, re.BaseRelation)
			for _, e := range re.Owns {
				BaseEntityList = append(BaseEntityList, e)
				// chunkSet.Add(e.ChunkId)
			}
			for _, e := range re.Owes {
				BaseEntityList = append(BaseEntityList, e)
				// chunkSet.Add(e.ChunkId)
			}
		}
	}

	return &BaseEntityList, &BaseRelationList, nil
}

func (b *RagBuilder) queryEntityVector(content string, docId string) ([]*ragStore.DocVector, error) {
	docs, err := b.vectordb.QueryEntity(content, docId, b.doEmbedding)
	return docs, err
}

func (b *RagBuilder) queryRelationVector(content string, docId string) ([]*ragStore.DocVector, error) {

	docs, err := b.vectordb.QueryRelation(content, docId, b.doEmbedding)
	return docs, err
}

func (b *RagBuilder) getChunkVector(chkSet *EntitySet, docId string) ([]*ragStore.ChunkVector, error) {
	if chkSet.Cardinality() < 1 {
		return nil, nil
	}
	var ids []string
	chkSet.Each(func(ckid string) bool {
		ids = append(ids, ckid)
		return true
	})

	docs, err := b.vectordb.GetChunk(docId, ids)
	return docs, err
}

func (b *RagBuilder) queryEntityGraph(docs []*ragStore.DocVector, docId string) (*EntityList, error) {
	entityList := make(EntityList, 0)
	for _, doc := range docs {
		entities, err := b.graphdb.GetNodeByName(doc.SrcName, docId)
		if err != nil {
			return nil, err
		}

		for _, entity := range *entities {
			entity.Rank = len(entity.Owns)
		}
		entityList = append(entityList, *entities...)
		// TODO: Query Entity Graph
	}
	return &entityList, nil
}

func (b *RagBuilder) queryEntityReverseGraph(docs []*ragStore.DocVector, docId string) (*EntityReverseList, error) {
	entityList := make(EntityReverseList, 0)
	for _, doc := range docs {
		entities, err := b.graphdb.GetReverseNode(doc.SrcName, docId)
		if err != nil {
			return nil, err
		}

		if len(*entities) < 1 {
			continue
		}

		for _, entity := range *entities {
			entity.Rank = len(entity.Owes)
		}
		entityList = append(entityList, *entities...)
	}
	return &entityList, nil
}

func (b *RagBuilder) queryRelationGraph(docs []*ragStore.DocVector, docId string) (*RelationEntityList, error) {
	relationList := make(RelationEntityList, 0)
	for _, doc := range docs {
		relations, err := b.graphdb.GetRelationEntity(doc.SrcName, doc.TgtName, docId)
		if err != nil {
			return nil, err
		}

		if len(*relations) < 1 {
			continue
		}

		relationList = append(relationList, *relations...)
	}

	return &relationList, nil
}

func (b *RagBuilder) buildSystemPromt(ragContent *RagContent, useDefault bool) (string, error) {
	if ragContent == nil {
		return "", errors.New("ragContent is nil")
	}
	type Prompt struct {
		Relations    string
		Entities     string
		Sources      string
		ResponseType string
	}

	Prompts := Prompt{
		Entities:     ragContent.Entities,
		Relations:    ragContent.Relationships,
		Sources:      ragContent.Sources,
		ResponseType: "Multiple Paragraphs",
	}

	tmp := RAG_RESPONSE
	if useDefault == false {
		tmp = RAG_ADDITION_PROMPT
	}
	tmpl, err := template.New("prompt").Parse(tmp)
	if err != nil {
		return "", err
	}
	buf := bytes.Buffer{}
	err = tmpl.Execute(&buf, Prompts)
	if err != nil {
		return "", err
	}

	return buf.String(), nil
}

func (b *RagBuilder) buildNaiveSystemPromt(naiveTxt *string, content *string, useDefault bool) (string, error) {
	type Prompt struct {
		ContextData  string
		ResponseType string
		// Question     string
	}

	Prompts := Prompt{
		ContextData:  *naiveTxt,
		ResponseType: "Multiple Paragraphs",
		// Question:     *content,
	}

	tmp := NAVIE_RAG_RESPONSE
	if useDefault == false {
		tmp = NAVIE_RAG_ADDITION_PROMPT
	}
	tmpl, err := template.New("prompt").Parse(tmp)
	if err != nil {
		return "", err
	}
	buf := bytes.Buffer{}
	err = tmpl.Execute(&buf, Prompts)
	if err != nil {
		return "", err
	}

	return buf.String(), nil
}

func (b *RagBuilder) buildKeyWordPromt(content string) (string, error) {
	type Prompt struct {
		Query string
	}

	Prompts := Prompt{
		Query: content,
	}

	tmpl, err := template.New("prompt").Parse(KEYWORDS_EXTRACTION)
	if err != nil {
		return "", err
	}

	buf := bytes.Buffer{}
	err = tmpl.Execute(&buf, Prompts)
	if err != nil {
		return "", err
	}

	return buf.String(), nil
}

func (b *RagBuilder) QueryLLM(ctx context.Context, content string, system string) (string, error) {
	msgs := []llm.RoleContent{}

	if system != "" {
		msgs = append(msgs, llm.RoleContent{Role: llm.RoleSys, Content: system})

	}

	msgs = append(msgs, llm.RoleContent{Role: llm.RoleUser, Content: content})

	req := &llm.CompleteReq{Msgs: msgs}
	rsp, err := b.query.Complete(ctx, req)
	if err != nil {
		return "", err
	}
	if rsp == nil {
		return "", errors.New("rsp is nil")
	}
	return rsp.GetContent(), nil
}

func (b *RagBuilder) QueryLLMBasic(ctx context.Context, prompt string) (string, error) {
	if len(prompt) < 1 {
		return "", errors.New("prompt is empty")
	}
	msgs := []llm.RoleContent{}

	msgs = append(msgs, llm.RoleContent{Role: llm.RoleUser, Content: prompt})
	Options := make(map[string]interface{})

	Options["temperature"] = 0.01
	Options["num_predict"] = 512

	req := &llm.CompleteReq{Msgs: msgs}
	rsp, err := b.query.Complete(ctx, req)
	if err != nil {
		return "", err
	}

	if rsp == nil {
		return "", errors.New("rsp is nil")
	}
	return rsp.GetContent(), nil
}

func (b *RagBuilder) QuerySummarTopic(docId string) ([]*ragStore.ChunkSummary, error) {
	return b.vectordb.QueryGSummary(docId)
}
